package edu.cuit.zhuyimeng.framework.common.util;

import cn.hutool.core.util.HashUtil;

import java.util.*;

/**
 * 一致性哈希工具
 */
public class ConsistentHash<T> {

    /**
     * 复制节点个数
     */
    private final int numberOfReplicas;

    private final Queue<T> excludedNodes = new LinkedList<>();

    /**
     * 一致性哈希环
     */
    private final SortedMap<Integer, T> circle = new TreeMap<>();

    public ConsistentHash(int numberOfReplicas, Collection<T> nodes) {
        this.numberOfReplicas = numberOfReplicas;
        for (T node : nodes) {
            add(node);
        }
    }

    public void add(T node) {
        for (int i = 0; i < numberOfReplicas; i++) {
            circle.put(HashUtil.fnvHash(node.toString() + i), node);
        }
    }

    /**
     * 移除节点的同时移除相应的虚拟节点
     * @param node 节点对象
     */
    public void remove(T node) {
        for (int i = 0; i < numberOfReplicas; i++) {
            circle.remove(HashUtil.fnvHash(node.toString() + i));
        }
    }

    /**
     * 获得一个最近的顺时针节点
     * @param key 为给定键取Hash，取得顺时针方向上最近的一个虚拟节点对应的实际节点
     * @return 节点对象
     */
    public T get(Object key) {
        if (circle.isEmpty()) {
            return null;
        }
        int hash = HashUtil.fnvHash(key.toString());
        if (!circle.containsKey(hash)) {
            SortedMap<Integer, T> tailMap = circle.tailMap(hash);
            hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
        }
        //正好命中
        return circle.get(hash);
    }

    /**
     * 排除元素（不计算）
     * @param node 节点对象
     */
    public void excludeNode(T node) {
        remove(node);
        excludedNodes.offer(node);
    }

    /**
     * 取消排除节点（最先排除的节点）
     */
    public T includeNode() {
        T node = excludedNodes.poll();
        if (node == null) {
            return null;
        }
        add(node);
        return node;
    }

    /**
     * 取消排除节点
     * @param node 节点对象
     */
    public void includeNode(T node) {
        excludedNodes.remove(node);
        add(node);
    }

}
